from typing import Callable, Optional, Union
import uuid
from pydantic import BaseModel
from sqlalchemy.orm import Session
from json import dumps, loads
from ...models import StateDBModel, State
from .utils import NestedModelEncoder


class StateCore:
    def __init__(self, make_session: Callable[[], Session], project_token: str):
        self.session = make_session
        self.project_token = project_token

    def __make_state_instance(self, data: Optional[str] = None, parent_state: Optional[str] = None):
        if data is None:
            data = dumps({
                'idea': '',
                'sources': [],
                'topics': [],
                'review': '',
            })
        new_id = str(uuid.uuid4())
        state = StateDBModel(
            id=new_id,
            data=data,
            parent_state=parent_state,
            project_token=self.project_token,
        )
        return state

    async def make_initial_state(self) -> State:
        state = self.__make_state_instance(parent_state=None)
        session = self.session()
        with session.begin():
            session.add(state)
        return State.model_validate(state)

    async def get_recent_state(self) -> State:
        session = self.session()
        with session.begin():
            current_state = session.query(StateDBModel).filter(
                StateDBModel.project_token == self.project_token
            ).order_by(StateDBModel.timestamp.desc()).first()
        if current_state is None:
            current_state = await self.make_initial_state()
        return State.model_validate(current_state)

    async def get_state(self, state_id: Optional[str]) -> Optional[State]:
        if state_id is None:
            return await self.get_recent_state()
        session = self.session()
        with session.begin():
            state = session.query(StateDBModel).filter(
                StateDBModel.project_token == self.project_token,
                StateDBModel.id == state_id
            ).first()
        return State.model_validate(state)

    async def get_data(self, state: State) -> dict:
        return loads(state.data)

    def __serialize_data(self, data: Union[dict, BaseModel]) -> str:
        if isinstance(data, BaseModel):
            return data.model_dump_json()
        return dumps(data, ensure_ascii=False, cls=NestedModelEncoder)

    async def append_state(self, parent_state: State, new_state_data: dict) -> State:
        data = self.__serialize_data(new_state_data)
        session = self.session()
        with session.begin():
            state = self.__make_state_instance(data=data, parent_state=parent_state.id)
            session.add(state)
        return State.model_validate(state)
